import random

# Pomocne funkcije
def mod_pow(a, n, m):
    result = 1
    a = a % m
    while n > 0:
        if n % 2 == 1:
            result = (result * a) % m
        a = (a * a) % m
        n = n // 2
        
    return result

def miller_rabin(n, k):
    if n <= 3:
        if n == 1:
            return False
        return True
  
    # n prost => n neparan => n = (2 ^ r) * d + 1
    d = n - 1
    r = 0
    while d % 2 == 0:
        r = r + 1
        d = d // 2
        
    for i in range(k):
        a = random.randrange(2, n - 1)
        
        x = mod_pow(a, d, n)

        if x == 1 or x == n - 1:
            continue
            
        wittness = True
        
        for j in range(r - 1):
            x = mod_pow(x, 2, n)
            if x == 1:
                return False
            if x == n - 1: # n - 1 = -1 (mod n)
                wittness = False
                break
        
        if wittness:
            return False
    return True


def get_prime(limit, k = 20):
    is_prime = False
    while not is_prime:
        n = random.randrange(limit)
        is_prime = miller_rabin(n, k)
    return n

# Pomoćna funkcija, prošireni Euklidov algoritam
def gcd(a, b):
    if b == 0:
        return a
    return gcd(b, a % b)

# Pomoćna funkcija, prošireni Euklidov algoritam
def ext_gcd(a, b):
    if b == 0:
        return (a, 1, 0)
    g, x, y = ext_gcd(b, a % b)
    return (g, y, x - a // b * y)

def mod_inv(a, m):
    g, x, y = ext_gcd(a, m)
    if g != 1:
        print("Vrednosti a i m nisu uzajamno proste!")
    else:
        return x % m

class MasseyOmura:
    def __init__(self, q):
        self.q = q
        while True:
            a = random.randrange(2, q-1)
            if gcd(a, q-1) == 1:
                self.a = a
                break
                
        self.a_inv = mod_inv(a, q-1)
        
    def encrypt(self, key):
        return pow(key, self.a, self.q)
    
    def decrypt(self, key):
        return pow(key, self.a_inv, self.q)

q = get_prime(2**100)

Amo = MasseyOmura(q)
Bmo = MasseyOmura(q)

key = random.randrange(2**10)
print(f'Key: {key}')

key_ph1A = Amo.encrypt(key)
print(f'A ---[{key_ph1A}]--> B')

key_ph1B = Bmo.encrypt(key_ph1A)
print(f'A <--[{key_ph1B}]--- B')

key_ph2A = Amo.decrypt(key_ph1B)
print(f'A ---[{key_ph2A}]--> B')

key_dec = Bmo.decrypt(key_ph2A)
print(f'B decrypts key: {key_dec}')

Key: 866
A ---[430472745990960853934299964940]--> B
A <--[25402872483233575461560422833]--- B
A ---[479643927642481006654566144552]--> B
B decrypts key: 866